from aimodel import embedding_model

import faiss
from langchain_community.vectorstores import FAISS
from langchain_community.docstore.in_memory import InMemoryDocstore

index = faiss.IndexFlatL2(len(embedding_model.embed_query("hello world")))

faiss_embedding = FAISS(
    embedding_function=embedding_model,
    index=index,
    docstore= InMemoryDocstore(),
    index_to_docstore_id={}
)

if __name__ == "__main__":
    from to_document import to_document
    import json
    with open(r"D:\kelun\文档\B线\B线氯化钠注射液工艺规程.json", "r", encoding="utf-8") as f:
        data = json.load(f)
    docus = to_document(data)
    faiss_embedding.add_documents(docus)
    retr = faiss_embedding.as_retriever()
    res = retr.invoke("氯化钠")
    print(res)
    for re in res:
        print(re.page_content)
        print(re.metadata)
        print(re.metadata["index"])